# -*- coding: utf-8 -*-
# @Function: Excel文件读取器
# @Description: 负责读取和解析Excel文件中的问答对数据，支持多种格式的Excel文件
# @Usage: 被importer.py导入使用，处理Excel文件读取和解析

import os
import pandas as pd
from typing import List, Tuple

class ExcelReader:
    def __init__(self):
        """初始化Excel读取器"""
        pass

    def read_file(self, file_path: str) -> pd.DataFrame:
        """读取Excel文件

        Args:
            file_path (str): Excel文件路径

        Returns:
            pd.DataFrame: Excel数据
        """
        try:
            df = pd.read_excel(file_path)
            print(f"成功读取Excel文件: {os.path.basename(file_path)}, 共 {len(df)} 行数据")
            print(f"列名: {list(df.columns)}")
            return df
        except Exception as e:
            print(f"读取Excel文件失败: {e}")
            raise

    def parse_qa_data(self, df: pd.DataFrame) -> List[Tuple[str, str, str]]:
        """解析Excel中的QA数据

        Args:
            df (pd.DataFrame): Excel数据

        Returns:
            List[Tuple[str, str, str]]: 问题、答案和上下文的列表
        """
        qa_pairs = []

        try:
            # 假设Excel文件有'question'、'answer'和'context'列，或者类似的列名
            question_cols = ['question', '问题', '提问', 'Question', 'Q']
            answer_cols = ['answer', '答案', 'Answer', '回答', 'A']
            context_cols = ['context', '上下文', '背景']

            question_col = None
            answer_col = None
            context_col = None

            # 查找问题列
            for col in question_cols:
                if col in df.columns:
                    question_col = col
                    break

            # 查找答案列
            for col in answer_cols:
                if col in df.columns:
                    answer_col = col
                    break

            # 查找上下文列 (可选)
            for col in context_cols:
                if col in df.columns:
                    context_col = col
                    break

            if not question_col or not answer_col:
                print(f"未找到问题或答案列。可用列: {list(df.columns)}")
                return qa_pairs

            print(f"使用问题列: {question_col}, 答案列: {answer_col}, 上下文列: {context_col}")

            # 提取QA对
            for index, row in df.iterrows():
                question = str(row[question_col]).strip()
                answer = str(row[answer_col]).strip()
                context = str(row[context_col]).strip() if context_col and pd.notna(row[context_col]) else None

                # 过滤空值和无效数据
                if question and answer and question != 'nan' and answer != 'nan':
                    qa_pairs.append((question, answer, context))

            print(f"解析到 {len(qa_pairs)} 个有效问答对")
            return qa_pairs

        except Exception as e:
            print(f"解析Excel QA数据失败: {e}")
            return qa_pairs

    def get_excel_files(self, folder_path: str) -> List[str]:
        """获取文件夹中的所有Excel文件

        Args:
            folder_path (str): 文件夹路径

        Returns:
            List[str]: Excel文件路径列表
        """
        if not os.path.exists(folder_path):
            print(f"文件夹不存在: {folder_path}")
            return []

        excel_files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) 
                      if f.endswith('.xlsx')]

        if not excel_files:
            print(f"在 {folder_path} 中没有找到任何Excel文件")
        else:
            print(f"找到 {len(excel_files)} 个Excel文件")

        return excel_files

    def extract_reservoir_name(self, filename: str) -> str:
        """从文件名提取水库名称

        Args:
            filename (str): 文件名（包含.xlsx扩展名）

        Returns:
            str: 水库名称（不含扩展名）
        """
        return os.path.splitext(os.path.basename(filename))[0]